import asyncio
from edge_tts import Communicate
import os
from flask import Flask, render_template, request, jsonify, send_from_directory
import subprocess
import uuid

app = Flask(__name__)
UPLOAD_FOLDER = 'static/output'
STATIC_DIR = os.path.join(os.path.dirname(__file__), "static")
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

MALE_VOICES = {
    "zh-CN-YunjianNeural", "zh-CN-YunxiNeural", "zh-CN-YunxiaNeural",
    "zh-CN-YunyangNeural", "zh-HK-WanLungNeural", "zh-TW-YunJheNeural",
    "zu-ZA-ThembaNeural"
}


async def text_to_speech(text, voice, output_file):
    communicate = Communicate(text=text, voice=voice)
    await communicate.save(output_file)

# 生成视频
def generate_video(face_image, audio_file, output_file):
    os.system(f"python /root/Easy-Wav2Lip/inference.py \
    --checkpoint_path /root/autodl-tmp/checkpoints/Wav2Lip_GAN.pth \
    --face {face_image} \
    --audio {audio_file}")
    os.system('cp /root/autodl-tmp/final_out/final.mp4 /root/tts_project/static/output/final.mp4')


@app.route('/')
def index():
    return render_template('index.html')

@app.route('/generate', methods=['POST'])
def generate():
    try:
        # 获取前端数据
        data = request.json
        text = data.get('text', '')
        voice = data.get('voice', 'zh-CN-XiaoxiaoNeural')
        
        audio_file = os.path.join(UPLOAD_FOLDER, 'gen.wav')
        video_file = os.path.join(UPLOAD_FOLDER, 'final.mp4')
        
        face_image = '/root/tts_project/static/woman.jpg'
        if voice in MALE_VOICES:                                  # ② 默认男头像
            face_image = os.path.join(STATIC_DIR, "man.jpg")
        else:                                                       # ③ 默认女头像
            face_image = os.path.join(STATIC_DIR, "woman.jpg")
        
        # 执行异步任务
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(text_to_speech(text, voice, audio_file))
        loop.close()
        
        # 生成视频
        generate_video(face_image, audio_file, video_file)
        
        return jsonify({
            'success': True,
            'video_url': f'/static/output/final.mp4',
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

@app.route('/static/output/<filename>')
def serve_video(filename):
    return send_from_directory(UPLOAD_FOLDER, filename)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=6006, debug=True)